import mlflow
import pandas as pd

def generate_recommendations_with_comparision(
    experiment_ids,
    aggregation_function="common_features",
    main_note="sizes_acts",
    note_to_compare="sizes_L2_without_acts",
    group_type="sim"
):
    all_rows = []

    for exp_id in experiment_ids:
        runs = mlflow.search_runs(
            experiment_ids=[exp_id],
            output_format="list"
        )

        for run in runs:
            note = run.data.params.get("note")
            fusion = run.data.params.get("SAE_fusion_strategy")
            group = run.data.params.get("group_type")

            if fusion != aggregation_function or group != group_type:
                continue

            dataset = run.data.params.get("dataset", f"Exp-{exp_id}")
            dim = int(run.data.params.get("embedding_dim", 0))
            topk = int(run.data.params.get("top_k", 0))

            row_key = (dataset, dim, topk)
            model_type = "main" if note == main_note else "compare" if note == note_to_compare else None
            if not model_type:
                continue

            metrics = {
                ("G/mean", model_type): run.data.metrics.get("CommonItemsNDCG20/median"),
                ("U/mean", model_type): run.data.metrics.get("NDCG20/mean"),
                ("U/min", model_type): run.data.metrics.get("NDCG20/min"),
                ("Pop", model_type): run.data.metrics.get("Popularity/mean"),
            }

            all_rows.append((row_key, metrics))

    # Combine metrics into a dictionary
    records = {}
    for key, metrics in all_rows:
        if key not in records:
            records[key] = {}
        records[key].update(metrics)

    df = pd.DataFrame.from_dict(records, orient="index")
    df.index.names = ["Dataset", "Dimensions", "TopK"]

    # Calculate % differences
    result_cols = []
    for metric in sorted(set(k[0] for k in df.columns)):
        main_col = (metric, "main")
        compare_col = (metric, "compare")
        percent_col = (metric, "% change")

        if main_col in df.columns and compare_col in df.columns:
            df[percent_col] = ((df[main_col] - df[compare_col]) / df[compare_col].abs()) * 100
            result_cols.extend([main_col, percent_col, compare_col])
        elif main_col in df.columns:
            result_cols.append(main_col)

    # Keep only main and percent change columns, sort them by metric
    df = df[result_cols]
    df = df.round(2)
    df = df.sort_index(axis=1, level=0).sort_values(by=["Dataset", "Dimensions", "TopK"])
    

    return df.reset_index()

import mlflow
import pandas as pd

def generate_common_features_table(experiment_ids, note="sizes_acts"):
    all_rows = []

    for exp_id in experiment_ids:
        runs = mlflow.search_runs(
            experiment_ids=[exp_id],
            filter_string=f"params.note = '{note}'",
            output_format="list"
        )

        for run in runs:
            params = run.data.params
            metrics = run.data.metrics
            
            if params.get("note") != note:
                continue
            

            dataset = params.get("dataset", f"Exp-{exp_id}")
            group_type = params.get("group_type", "unknown").capitalize()
            dim = int(params.get("embedding_dim", 0))
            topk = int(params.get("top_k", 0))
            value = metrics.get("common_features/mean", None)

            if value is None:
                continue

            row_key = (dim, topk)
            all_rows.append((row_key, (dataset, group_type), value))

    # Create dictionary for DataFrame
    records = {}
    for row_key, col_key, value in all_rows:
        if row_key not in records:
            records[row_key] = {}
        records[row_key][col_key] = value

    df = pd.DataFrame.from_dict(records, orient="index")
    df.index.names = ["Dimensions", "TopK"]

    # Reorder columns
    if not df.empty:
        datasets = ["MovieLens",'LastFM1k']
        subcols = ["Sim", "Random", "Div"]
        col_order = [(d, s) for d in datasets for s in subcols]
        df = df.reindex(columns=pd.MultiIndex.from_tuples(col_order))
        

    return df.reset_index()

experiments = ['962723054918039068', '228719589483846826']
generate_common_features_table(experiments, note="sizes_L2_without_acts")

experiment_ids = ['333391697323445885', '523100174176986081']
table = generate_recommendations_with_comparision(
    experiment_ids,
    aggregation_function="common_features",
    main_note="sizes_L2_without_acts",
    note_to_compare="sizes_L2_with_acts",
    group_type="sim"
)
table

table.groupby(["Dataset"]).mean()

table.groupby(["Dataset"]).std()

table.groupby(["Dataset"]).max()

table.groupby(["Dataset"]).min()

experiment_ids = ['333391697323445885', '523100174176986081']
table = generate_recommendations_with_comparision(
    experiment_ids,
    aggregation_function="average",
    main_note="sizes_L2_without_acts",
    note_to_compare="sizes_L2_with_acts",
    group_type="sim"
)
table

table.groupby(["Dataset"]).mean()

table.groupby(["Dataset"]).std()

table.groupby(["Dataset"]).max()

table.groupby(["Dataset"]).min()

experiment_ids = ['333391697323445885', '523100174176986081']
generate_recommendations_with_comparision(
    experiment_ids,
    aggregation_function="average",
    main_note="sizes_L2_without_acts",
    note_to_compare="sizes_L2_with_acts",
    group_type="random"
)

experiment_ids = ['333391697323445885', '523100174176986081']
generate_recommendations_with_comparision(
    experiment_ids,
    aggregation_function="average",
    main_note="sizes_L2_without_acts",
    note_to_compare="sizes_L2_with_acts",
    group_type="div"
)

experiment_ids = ['333391697323445885', '523100174176986081']
generate_recommendations_with_comparision(
    experiment_ids,
    aggregation_function="common_features",
    main_note="sizes_L2_without_acts",
    note_to_compare="sizes_L2_with_acts",
    group_type="random"
)

experiment_ids = ['333391697323445885', '523100174176986081']
generate_recommendations_with_comparision(
    experiment_ids,
    aggregation_function="common_features",
    main_note="sizes_L2_without_acts",
    note_to_compare="sizes_L2_with_acts",
    group_type="div"
)